"""
从 coco2014 数据集中随机采样固定数量的 captions
"""
import sys
sys.path.append('.')  # 为了能够直接在根目录下运行 debug
from tqdm.auto import tqdm
from loguru import logger
from reflow.utils import set_seed
import numpy as np
import json
from argparse import Namespace, ArgumentParser
# import random


def prepare_args():
    parser = ArgumentParser()
    parser.add_argument(
        "--split",
        type=str,
    )
    parser.add_argument(
        "--nums",
        type=int,
    )
    parser.add_argument(
        "--file_path",
        type=str,
    )
    parser.add_argument(
        "--phase",
        type=str,
        default="2014",
        help="2014 or 2017"
    )
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = prepare_args()

    # ! debug or convenience
    ######################################
    args.split = 'val'
    args.nums = 5000
    args.file_path = 'data/reflow/coco20147_val_random5k.txt'
    args.seed = 1891
    args.phase = "2017"
    ######################################

    set_seed(args.seed)
    assert args.split in ['train', 'val']

    json_path = f'data/coco{args.phase}/annotations/captions_{args.split}{args.phase}.json'
    logger.info(f'loading json file from <<{json_path}>> ...')
    json_info = json.load(open(json_path, 'r'))
    logger.info(f'done!')

    logger.info(f'merge all captions to a list...')
    all_caps = []
    for a in tqdm(json_info['annotations']):
        cap = a['caption']
        # ! coco 中有的文本格式错误需要做处理
        cap = cap.replace('\n', '').strip(' ')
        all_caps.append(cap)
    logger.info(f'done!')

    logger.info(f'randomly sampling captions ...')
    all_caps = np.array(all_caps)
    rand_indices = np.random.randint(low=0, high=len(all_caps), size=args.nums)
    rand_caps = all_caps[rand_indices].tolist()
    logger.info(f'done')

    logger.info(f'writing caption samples to {args.file_path} ...')
    with open(args.file_path, 'w') as f:
        f.write('\n'.join(rand_caps))
    logger.info('done!')
